"""
Utility functions for Dumbo benchmark
"""

import os
import sys
from typing import Any, Dict


class BenchError(Exception):
    """Custom exception for benchmark errors"""
    pass


class Print:
    """Colored print utilities"""
    
    @staticmethod
    def info(msg: str, **kwargs):
        """Print info message"""
        print(f"[INFO] {msg}", **kwargs)
    
    @staticmethod
    def success(msg: str, **kwargs):
        """Print success message"""
        print(f"[SUCCESS] {msg}", **kwargs)
    
    @staticmethod
    def warning(msg: str, **kwargs):
        """Print warning message"""
        print(f"[WARNING] {msg}", **kwargs)
    
    @staticmethod
    def error(msg: str, **kwargs):
        """Print error message"""
        print(f"[ERROR] {msg}", **kwargs)
    
    @staticmethod
    def heading(msg: str, **kwargs):
        """Print heading message"""
        print(f"\n{'='*60}")
        print(f"{msg}")
        print(f"{'='*60}", **kwargs)
    
    @staticmethod
    def subheading(msg: str, **kwargs):
        """Print subheading message"""
        print(f"\n{'-'*40}")
        print(f"{msg}")
        print(f"{'-'*40}", **kwargs)


def load_config(config_path: str) -> Dict[str, Any]:
    """Load configuration from JSON or YAML file"""
    import json
    import yaml
    
    if not os.path.exists(config_path):
        raise BenchError(f"Configuration file not found: {config_path}")
    
    with open(config_path, 'r') as f:
        if config_path.endswith('.json'):
            return json.load(f)
        elif config_path.endswith(('.yaml', '.yml')):
            return yaml.safe_load(f)
        else:
            raise BenchError(f"Unsupported configuration file format: {config_path}")


def save_config(config: Dict[str, Any], config_path: str):
    """Save configuration to JSON or YAML file"""
    import json
    import yaml
    
    os.makedirs(os.path.dirname(config_path), exist_ok=True)
    
    with open(config_path, 'w') as f:
        if config_path.endswith('.json'):
            json.dump(config, f, indent=2)
        elif config_path.endswith(('.yaml', '.yml')):
            yaml.dump(config, f, default_flow_style=False)
        else:
            raise BenchError(f"Unsupported configuration file format: {config_path}")


def ensure_directory(path: str):
    """Ensure directory exists"""
    os.makedirs(path, exist_ok=True)


def get_timestamp() -> str:
    """Get current timestamp string"""
    from datetime import datetime
    return datetime.now().strftime("%Y%m%d_%H%M%S")


def run_command(cmd: str, cwd: str = None) -> tuple:
    """Run shell command and return (returncode, stdout, stderr)"""
    import subprocess
    
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            cwd=cwd,
            capture_output=True,
            text=True,
            timeout=300  # 5 minute timeout
        )
        return result.returncode, result.stdout, result.stderr
    except subprocess.TimeoutExpired:
        return -1, "", "Command timed out"
    except Exception as e:
        return -1, "", str(e)


def check_dependencies():
    """Check if required dependencies are available"""
    dependencies = {
        'go': 'go version',
        'python3': 'python3 --version',
        'tmux': 'tmux -V'
    }
    
    missing = []
    for dep, cmd in dependencies.items():
        returncode, _, _ = run_command(cmd)
        if returncode != 0:
            missing.append(dep)
    
    if missing:
        raise BenchError(f"Missing dependencies: {', '.join(missing)}")
    
    Print.success("All dependencies are available")


def validate_config(config: Dict[str, Any], required_keys: list):
    """Validate configuration has required keys"""
    missing_keys = []
    for key in required_keys:
        if key not in config:
            missing_keys.append(key)
    
    if missing_keys:
        raise BenchError(f"Configuration missing required keys: {', '.join(missing_keys)}")


def format_bytes(bytes_value: int) -> str:
    """Format bytes to human readable string"""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if bytes_value < 1024.0:
            return f"{bytes_value:.1f} {unit}"
        bytes_value /= 1024.0
    return f"{bytes_value:.1f} PB"


def format_duration(seconds: float) -> str:
    """Format duration to human readable string"""
    if seconds < 60:
        return f"{seconds:.1f}s"
    elif seconds < 3600:
        return f"{seconds/60:.1f}m"
    else:
        return f"{seconds/3600:.1f}h"

